import pymysql
import socket
import datetime
import os
import json
import csv
# ======================== 自定义参数配置区 ======================== #
# 数据库连接配置
DB_CONFIG = {
    'host': '172.16.30.54',
    'port': 3306,
    'user': 'root',
    'password': 'IPS2025a%',
    'database': 'iip_numerical_algorithm_mini_program',
    'cursorclass': pymysql.cursors.DictCursor
}

# 查询参数配置
QUERY_CLICK_PARAMS = {
    'event_type': 'click',
    'page_ids': ['informationInfo', 'courseInfo', 'videoInfo', 'recommendInfo'],
    'table_name': 'tracking_data',  # 请替换为实际表名
    'timeout': 5  # 网络连接超时时间（秒）
}
# 查询参数配置
QUERY_VIEWPORT_PARAMS = {
    'event_type': 'viewport',
    'page_ids': ['informationInfo', 'courseInfo', 'videoInfo', 'recommendInfo'],
    'table_name': 'tracking_data',  # 请替换为实际表名
    'timeout': 5  # 网络连接超时时间（秒）
}
# ======================== 自定义参数配置区结束 ======================== #

def test_connection(host, port, timeout=5):
    """测试网络连接"""
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        result = sock.connect_ex((host, port))
        if result == 0:
            print(f"成功连接到 {host}:{port}")
        else:
            print(f"无法连接到 {host}:{port}，错误码: {result}")
        sock.close()
    except Exception as e:
        print(f"测试连接时出错: {e}")

def query_click_events():
    """查询符合条件的点击事件"""
    try:
        # 使用配置的数据库连接信息
        with pymysql.connect(**DB_CONFIG) as connection:
            with connection.cursor() as cursor:
                # 准备查询参数
                page_ids = QUERY_CLICK_PARAMS['page_ids']
                placeholders = ', '.join(['%s'] * len(page_ids))
                
                # 执行SQL查询
                sql = f"SELECT * FROM {QUERY_CLICK_PARAMS['table_name']} WHERE event_type = %s AND page_id IN ({placeholders})"
                cursor.execute(sql, [QUERY_CLICK_PARAMS['event_type']] + page_ids)
                results = cursor.fetchall()
                
                return results
    
    except pymysql.Error as e:
        print(f"点击事件数据库操作出错: {e}")
        return None

def query_viewport_events():
    """查询符合条件的曝光事件"""
    try:
        # 使用配置的数据库连接信息
        with pymysql.connect(**DB_CONFIG) as connection:
            with connection.cursor() as cursor:
                # 准备查询参数
                page_ids = QUERY_VIEWPORT_PARAMS['page_ids']
                placeholders = ', '.join(['%s'] * len(page_ids))
                
                # 执行SQL查询
                sql = f"SELECT * FROM {QUERY_VIEWPORT_PARAMS['table_name']} WHERE event_type = %s AND page_id IN ({placeholders})"
                cursor.execute(sql, [QUERY_VIEWPORT_PARAMS['event_type']] + page_ids)
                results = cursor.fetchall()
                
                return results
    
    except pymysql.Error as e:
        print(f"曝光事件数据库操作出错: {e}")
        return None


# 以下为 match_events_and_label 函数定义
def match_events_and_label(click_events, viewport_events):
    """匹配点击事件和曝光事件，并为曝光事件添加label标记
    
    参数:
        click_events: 点击事件列表
        viewport_events: 曝光事件列表
    
    返回:
        带label标记的曝光事件列表，label=1表示匹配(点击事件在曝光事件5分钟内)，label=0表示未匹配
    """
    if not click_events or not viewport_events:
        print("未找到点击事件或曝光事件数据")
        return []
    
    # 预处理事件：删除指定列并解析extra_param
    click_events = process_events(click_events)
    viewport_events = process_events(viewport_events)
    
    # 为了提高匹配效率，按(openid, deviceid, page_id, content_id)分组存储点击事件
    click_events_grouped = {}    
    for click in click_events:
        # 获取内容标识(articleId或videoId)
        content_id = click['articleId'] if click['articleId'] != -1 else click['videoId']
        key = (click['openid'], click['deviceid'], click['page_id'], content_id)
        if key not in click_events_grouped:
            click_events_grouped[key] = []
        # 解析点击事件时间
        try:
            if isinstance(click['event_time'], datetime.datetime):
                click_time = click['event_time']
            else:
                click_time = datetime.datetime.strptime(click['event_time'], '%Y-%m-%d %H:%M:%S')
            click_events_grouped[key].append({'time': click_time, 'event': click})
        except ValueError:
            print(f"无法解析点击事件时间: {click['event_time']}")
            continue
    
    # 为每个曝光事件添加label标记
    labeled_viewport_events = []
    for viewport in viewport_events:
        # 获取内容标识(articleId或videoId)
        content_id = viewport['articleId'] if viewport['articleId'] != -1 else viewport['videoId']
        key = (viewport['openid'], viewport['deviceid'], viewport['page_id'], content_id)
        label = 0
        
        if key in click_events_grouped:
            try:
                if isinstance(viewport['event_time'], datetime.datetime):
                    viewport_time = viewport['event_time']
                else:
                    viewport_time = datetime.datetime.strptime(viewport['event_time'], '%Y-%m-%d %H:%M:%S')
                
                for click in click_events_grouped[key]:
                    time_diff = abs((click['time'] - viewport_time).total_seconds())
                    if time_diff <= 300:  # 5分钟内
                        label = 1
                        break
            except ValueError:
                print(f"无法解析曝光事件时间: {viewport['event_time']}")
        
        labeled_event = {'label': label}
        labeled_event.update(viewport)
        labeled_viewport_events.append(labeled_event)
    
    matched_count = sum(1 for event in labeled_viewport_events if event['label'] == 1)
    total_count = len(labeled_viewport_events)
    print(f"匹配统计: 总曝光事件数={total_count}, 匹配点击事件数={matched_count}, 匹配率={matched_count/total_count:.2%}")
    
    return labeled_viewport_events

# 添加新的process_events函数
def process_events(events):
    """预处理事件数据
    
    参数:
        events: 事件列表
    
    返回:
        预处理后的事件列表
    """
    processed_events = []
    for event in events:
        processed_event = event.copy()
        
        # 删除指定列
        columns_to_delete = ['user_id', 'device_id', 'id']
        for col in columns_to_delete:
            if col in processed_event:
                del processed_event[col]
        
        # 处理tenant_id为空的情况
        processed_event['tenant_id'] = processed_event.get('tenant_id') if processed_event.get('tenant_id') is not None else -1
        
        # 解析extra_param字段
        if 'extra_param' in processed_event:
            try:
                extra_param = json.loads(processed_event['extra_param'])
                # 提取所需字段
                openid_value = extra_param.get('openid', '')
                processed_event['openid'] = -1 if openid_value == '' else openid_value
                processed_event['deviceid'] = extra_param.get('deviceid') if extra_param.get('deviceid') is not None else -1
                processed_event['articleId'] = extra_param.get('articleId', -1)
                processed_event['videoId'] = extra_param.get('videoId', -1)
                processed_event['indexId'] = extra_param.get('indexId', -1)  # 新增提取indexId字段
                
                # 移除原始extra_param字段
                del processed_event['extra_param']
            except json.JSONDecodeError:
                print(f"警告: 无法解析extra_param字段: {processed_event.get('extra_param', '')}")
                # 设置默认值
                processed_event['openid'] = -1
                processed_event['deviceid'] = -1
                processed_event['articleId'] = -1
                processed_event['videoId'] = -1
                processed_event['indexId'] = -1  # 解析失败时设置默认值
                del processed_event['extra_param']
        else:
            # 如果没有extra_param字段，设置默认值
            processed_event['openid'] = -1
            processed_event['deviceid'] = -1
            processed_event['articleId'] = -1
            processed_event['videoId'] = -1
            processed_event['indexId'] = -1  # 无extra_param时设置默认值
        
        processed_events.append(processed_event)
    
    return processed_events

# 修改save_labeled_events_to_csv函数
def save_labeled_events_to_csv(labeled_events, output_dir):
    """将带标记的事件数据按天级别分桶保存为CSV文件
    
    参数:
        labeled_events: 带标记的事件列表
        output_dir: 输出目录路径
    """
    if not labeled_events:
        print("没有可保存的标记事件数据")
        return
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 按日期分组事件
    events_by_day = {}    
    for event in labeled_events:
        try:
            # 确保event_time是datetime对象
            if not isinstance(event['event_time'], datetime.datetime):
                event_time = datetime.datetime.strptime(event['event_time'], '%Y-%m-%d %H:%M:%S')
            else:
                event_time = event['event_time']
            
            # 提取日期作为分组键
            date_key = event_time.strftime('%Y-%m-%d')
            
            if date_key not in events_by_day:
                events_by_day[date_key] = []
            events_by_day[date_key].append(event)
        except Exception as e:
            print(f"处理事件时间时出错: {e}")
            continue
    
    # 为每个日期保存数据
    for date_key, events in events_by_day.items():
        # 创建日期子目录
        date_dir = os.path.join(output_dir, date_key)
        os.makedirs(date_dir, exist_ok=True)
        
        # 输出文件路径
        output_file = os.path.join(date_dir, 'labeled_viewport_events.csv')
        
        # 获取所有字段名，确保label是第一列
        fieldnames = ['label'] + [key for key in events[0].keys() if key != 'label']
        
        with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(events)
        print(f"{date_key}的数据已保存至: {output_file}")

# 执行测试和查询
if __name__ == "__main__":
    # 测试数据库连接
    test_connection(DB_CONFIG['host'], DB_CONFIG['port'], QUERY_CLICK_PARAMS['timeout'])
    # 执行查询
    click_events = query_click_events()
    viewport_events = query_viewport_events()
    
    # 匹配事件并添加label
    labeled_viewport_events = match_events_and_label(click_events, viewport_events)
    
    # 解析extra_param字段并按天保存到CSV文件
    output_dir = '/data/gongzhijia/data/rawdata'
    save_labeled_events_to_csv(labeled_viewport_events, output_dir)

# 做拆分